{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MSC 工具测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/media/pc/data/lxw/ai/tvm-book/doc/tutorials/msc\n"
     ]
    }
   ],
   "source": [
    "%cd ..\n",
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.contrib.msc.core.tools import ToolType"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'pruner'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ToolType.PRUNER"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "MSCTools 与 MSCGraph 协同工作，它们决定压缩策略并控制压缩过程。MSCTools 由 RuntimeManager 管理。\n",
    "\n",
    "```python\n",
    "from tvm.contrib.msc.core.transform import msc_transform\n",
    "from tvm.contrib.msc.core.runtime import create_runtime_manager\n",
    "from tvm.contrib.msc.core.tools import create_tool, MSC_TOOL\n",
    "\n",
    "# build runtime manager from module and mscgraphs\n",
    "optimized_mod, msc_graph, msc_config = msc_transform(mod, params)\n",
    "rt_manager = create_runtime_manager(optimized_mod, params, msc_config)\n",
    "\n",
    "# pruner is used for prune the model\n",
    "rt_manager.create_tool(MSC_TOOL.PRUNE, prune_config)\n",
    "\n",
    "# quantizer is used to do the calibration and quantize the model\n",
    "rt_manager.create_tool(MSC_TOOL.QUANTIZE, quantize_config)\n",
    "\n",
    "# collecter is used to collect the datas of each computational node\n",
    "rt_manager.create_tool(MSC_TOOL.COLLECT, collect_config)\n",
    "\n",
    "# distiller is used to do the knowledge distilliation\n",
    "rt_manager.create_tool(MSC_TOOL.DISTILL, distill_config)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MSCGym\n",
    "\n",
    "MSCGym 是 MSC 中自动压缩的平台。它的作用类似于 AutoTVM，但其架构更像 OpenAI-Gym。MSCGym 从压缩过程中提取任务，然后利用代理和环境之间的交互来为每个任务找到最佳行动。要使用 MSCGym 进行自动压缩，请为工具设置 `gym` 配置：\n",
    "\n",
    "```\n",
    "{\n",
    "      ...\n",
    "      \"quantize\": {\n",
    "        \"strategy_file\": \"msc_quantize.json\",\n",
    "        \"target\": \"tensorrt\",\n",
    "        “gym”:[\n",
    "          {\n",
    "            “record”:”searched_config.json”,\n",
    "            “env”:{\n",
    "              “strategy”:”distill_loss”\n",
    "            },\n",
    "            “agent”:{\n",
    "              “type”:”grid_search”,\n",
    "            }\n",
    "          },\n",
    "        ]\n",
    "      },\n",
    "      ...\n",
    "}\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "xxx",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
